import gymnasium as gym
import gymnasium.wrappers as gym_wrap

import DQN_model as DQN

env = gym.make("CarRacing-v2", continuous=False, render_mode="human")
env = DQN.SkipFrame(env, skip=4)
env = gym_wrap.GrayScaleObservation(env)
env = gym_wrap.ResizeObservation(env, shape=84)
env = gym_wrap.FrameStack(env, num_stack=4)
state, info = env.reset()
action_n = env.action_space.n
driver = DQN.Agent(
    state.shape,
    action_n,
    load_state='eval',
    load_model='DQN_740188.pt'
)
driver.epsilon = 0
episodes = 1
scores_array = []
timestep_arr = []
seeds_list = [i for i in range(10, 12, 1)]  # List of any seeds can be specified

for episode, sd in enumerate(seeds_list):
    state, info = env.reset(seed=sd)
    updating = True
    score = 0
    timestep = 0

    while updating:
        action = driver.take_action(state)
        state, reward, terminated, truncated, info = env.step(action)
        updating = not (terminated or truncated)
        score += reward
        timestep += 1
    scores_array.append(score)
    timestep_arr.append(timestep)
    print(f"Episode:{episode}, Score:{score:.2f}, Timesteps: {timestep}")

env.close()
